import glob
import json
import os
import random
import time

import cv2

'''
测试数据集精度统计
对比两者标注
labelmepose 和 pose_test_dataset100-labelme ()
'''
# from log import get_logger
from loguru import logger # 日志


def compare_annotation(test_config):
    def iou(box1, box2):
        """
        计算两个矩形框的 IOU

        参数:
        box1 -- 第一个矩形框 [x1, y1, x2, y2]
        box2 -- 第二个矩形框 [x1, y1, x2, y2]

        返回:
        iou -- IOU 比值 (float)
        """

        # 获取坐标
        x1_1, y1_1, x2_1, y2_1 = box1
        x1_2, y1_2, x2_2, y2_2 = box2

        # 计算交集区域的坐标
        inter_x1 = max(x1_1, x1_2)
        inter_y1 = max(y1_1, y1_2)
        inter_x2 = min(x2_1, x2_2)
        inter_y2 = min(y2_1, y2_2)

        # 如果没有交集
        if inter_x1 >= inter_x2 or inter_y1 >= inter_y2:
            return 0.0

        # 计算交集面积
        inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1)

        # 计算每个矩形框的面积
        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)

        # 并集面积
        union_area = area1 + area2 - inter_area

        # IOU
        iou_value = inter_area / union_area

        return iou_value

    # true_annotation_dir = test_config.get('true_annotation_dir')
    # pre_annotation_dir = test_config.get('pre_annotation_dir')

    cls_model_type = test_config.get('cls_model_type')
    is_save_img = test_config.get('is_save_img',False)
    save_dir = test_config.get('save_dir')
    img_glob = test_config.get('img_glob') # *.jpg

    logger.add(f'{save_dir}/log.txt', rotation='200KB')
    print = logger.info

    true_annotation_dir = img_glob[:-6]
    pre_annotation_dir = f'{save_dir}/annotated'
    # true_annotation_dir = r'D:\DATA\20250519RENBAO\caitu\pose_test_dataset250-labelme'  # pose_test_dataset100-labelme
    img_glob = f'{true_annotation_dir}/*.jpg'
    # pre_annotation_dir = r'D:\DATA\20250519RENBAO\temp\annotated'  # labelmepose
    save_dir = f'{save_dir}/val_show_img'

    test_cls_names = ['baggage_place', 'none_action']
    pre_cls_names = ['IP', 'NP']
    pre_tru_clsname_dict = {
        'ZhCMFangBao': 'baggage_place',
        'ZhCMFuKuang': 'baggage_place',
        'ZhCMNaKuang': 'none_action',
        'ZhCMZhanLi': 'none_action',
        'BMFangBao': 'none_action',
        'BMZhanLi': 'none_action',
        'AnJianYuan': 'none_action',
        # 'Other': 'none_action'
    }

    cls_colors = [ (255, 255, 255), (0, 0, 0)]  # 黑(不放包) 白(阳例, 放包)
    acc_cls_names = [0, 0]  # acc count
    # tru_pre_acc_matrix = {
    #     test_cls_names[0]: {
    #         pre_cls_names[0]: 0,
    #         pre_cls_names[1]: 0
    #     },
    #     test_cls_names[1]: {
    #         pre_cls_names[0]: 0,
    #         pre_cls_names[1]: 0
    #     }
    # } # true_pre
    tru_pre_acc_matrix = {}
    for tru_name in pre_tru_clsname_dict.values():
        tru_pre_acc_matrix[tru_name] = {}
        for pre_name in pre_tru_clsname_dict.keys():
            tru_pre_acc_matrix[tru_name][pre_name] = 0

    total_cls_names = [0, 0]  # total
    box_error_count = 0
    img_ls = glob.glob(img_glob)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    for img_path in img_ls:
        img_file_name = os.path.basename(img_path)  # '000000.jpg'
        img_name = os.path.splitext(img_file_name)[0]  # '000000'
        img = cv2.imread(img_path)
        pre_annotation_path = f"{pre_annotation_dir}/{img_name}.json"  # 根据文件名匹配
        true_annotation_path = f"{true_annotation_dir}/{img_name}.json"
        if not os.path.exists(pre_annotation_path) or not os.path.exists(true_annotation_path):
            continue

        pre_annotation = json.load(open(pre_annotation_path))
        true_annotation = json.load(open(true_annotation_path))

        for true_rect_shape_dict in true_annotation['shapes']:
            if not true_rect_shape_dict['shape_type'] == 'rectangle':
                continue
            [true_x1, true_y1], [true_x2, true_y2] = true_rect_shape_dict['points']
            true_cls_name = true_rect_shape_dict['label']
            true_x1, true_y1, true_x2, true_y2 = int(true_x1), int(true_y1), int(true_x2), int(true_y2)
            # 画 true 标注
            true_color = cls_colors[test_cls_names.index(true_cls_name)]
            cv2.rectangle(img, (true_x1, true_y1), (true_x2, true_y2), true_color, 2)  # bgr 白
            cv2.line(img, (true_x1, (true_y1+true_y2)//2), (true_x2, (true_y1+true_y2)//2), true_color, 2) # 横线

            # 找匹配的 pre_rect_shape_dict
            max_iou_value = 0
            max_iou_pre_rect_shape_dict_ind = 0
            for shape_dict_ind, pre_rect_shape_dict in enumerate(pre_annotation['shapes']):
                if not pre_rect_shape_dict['shape_type'] == 'rectangle':
                    continue
                [pre_x1, pre_y1], [pre_x2, pre_y2] = pre_rect_shape_dict['points']
                iou_value = iou([true_x1, true_y1, true_x2, true_y2], [pre_x1, pre_y1, pre_x2, pre_y2])
                if iou_value > max_iou_value:
                    max_iou_value = iou_value
                    max_iou_pre_rect_shape_dict_ind = shape_dict_ind

            is_error_box = max_iou_value < 0.1
            if is_error_box:  # 漏检
                box_error_count += 1
                print(f'loujian: {img_name}')
                continue
            max_iou_pre_rect_shape_dict = pre_annotation['shapes'][max_iou_pre_rect_shape_dict_ind]

            pre_cls_name = max_iou_pre_rect_shape_dict['label']
            [pre_x1, pre_y1], [pre_x2, pre_y2] = max_iou_pre_rect_shape_dict['points']
            pre_x1, pre_y1, pre_x2, pre_y2 = int(pre_x1), int(pre_y1), int(pre_x2), int(pre_y2)

            total_cls_names[test_cls_names.index(true_cls_name)] += 1
            # is_acc = test_cls_names.index(true_cls_name) == pre_cls_names.index(pre_cls_name)
            is_acc = true_cls_name == pre_tru_clsname_dict[pre_cls_name]
            if is_acc:
                acc_cls_names[test_cls_names.index(true_cls_name)] += 1

            tru_pre_acc_matrix[true_cls_name][pre_cls_name] += 1

            # 画结果图
            pre_color = cls_colors[test_cls_names.index(pre_tru_clsname_dict[pre_cls_name])]
            cv2.rectangle(img, (pre_x1, pre_y1), (pre_x2, pre_y2), pre_color, 6, lineType=cv2.LINE_4)  # 黑
            cv2.line(img, ((pre_x1+pre_x2)//2, pre_y1), ((pre_x1+pre_x2)//2, pre_y2), pre_color, 6)

            text_color = (0, 255, 0) if is_acc else (0, 0, 255)
            cv2.putText(img, f'{true_cls_name}/{pre_cls_name} {max_iou_value:.2f}', (true_x1, true_y1),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
            if is_save_img:
                cv2.imwrite(f'{save_dir}/{img_file_name}', img)

    print(f'img_glob {img_glob}')
    print(f'pre_annotation_dir {pre_annotation_dir}')
    print(f'true_annotation_dir {true_annotation_dir}')
    print(f'save_dir {save_dir}')
    print(f'cls_model_type {cls_model_type}')
    print(f'box_error_count {box_error_count}')
    for cls_name in test_cls_names:
        cls_ind = test_cls_names.index(cls_name)
        print(
            f'{cls_name} acc: {acc_cls_names[cls_ind]} / {total_cls_names[cls_ind]} {acc_cls_names[cls_ind] / total_cls_names[cls_ind]:.2f}')
    print(f'acc: {sum(acc_cls_names)} / {sum(total_cls_names)} {sum(acc_cls_names) / sum(total_cls_names):.2f}')

    # print(f'tru_pre_acc_matrix {tru_pre_acc_matrix}')
    # todo 画tru_pre_acc_matrix
    print(f'tru_pre_acc_matrix')
    for tru_name, pre_dict in tru_pre_acc_matrix.items():
        print(f'{tru_name} {pre_dict}')
    # tp = tru_pre_acc_matrix['baggage_place']['IP']
    # tn = tru_pre_acc_matrix['none_action']['NP']
    # fp = tru_pre_acc_matrix['none_action']['IP']
    # fn = tru_pre_acc_matrix['baggage_place']['NP']
    
    fuse_tru_pre_acc_matrix = {}
    for tru, pre_dict in tru_pre_acc_matrix.items():
        if not fuse_tru_pre_acc_matrix.get(tru):
            fuse_tru_pre_acc_matrix[tru] = {}
        for pre, count in pre_dict.items():
            fuse_pre = pre_tru_clsname_dict[pre]
            if not fuse_tru_pre_acc_matrix[tru].get(fuse_pre):
                fuse_tru_pre_acc_matrix[tru][fuse_pre] = 0
            fuse_tru_pre_acc_matrix[tru][fuse_pre] += count
    
    tp = fuse_tru_pre_acc_matrix['baggage_place']['baggage_place']
    fn = fuse_tru_pre_acc_matrix['baggage_place']['none_action']
    fp = fuse_tru_pre_acc_matrix['none_action']['baggage_place']

    tn = fuse_tru_pre_acc_matrix['none_action']['none_action']
    p = tp/(tp + fp)
    r = tp/(tp + fn)
    none_action_p = tn/(tn + fn)
    none_action_r = tn/(tn + fp)

    # p = tru_pre_acc_matrix['baggage_place']['IP']/(tru_pre_acc_matrix['baggage_place']['IP'] + tru_pre_acc_matrix['baggage_place']['NP']) #
    # r = tru_pre_acc_matrix['baggage_place']['IP']/(tru_pre_acc_matrix['baggage_place']['IP'] + tru_pre_acc_matrix['none_action']['IP'])
    print(f'none_action P: {none_action_p:.2f} R: {none_action_r:.2f}')
    print(f'baggage_place P: {p:.2f} R: {r:.2f}')
    # os.startfile(save_dir)

if __name__ == '__main__':
    # #
    # test_config = {
    #     'true_annotation_dir': r'D:\DATA\20250519RENBAO\caitu\pose_test_dataset250-labelme',
    #     # 'true_annotation_dir': r'D:\DATA\20250519RENBAO\caitu\20250603_pose_test_dataset250-labelme',
    #     'pre_annotation_dir': r'D:\DATA\20250519RENBAO\temp\annotated',
    #     'cls_model_type': r'img100_200_cls7_augmean3_pad',
    #     'is_save_img': False,
    # }

    root_path = r'D:\DATA\20250519RENBAO'
    cls_model_type = 'frontCamImg0_100_cls4_augMean3_padSquare'
    test_config = {
        'pose_model_path': rf"D:\CODE\ZXC\project_rbao\yolo11m-pose.pt",
        'pose_model_type': 'official',
        'pose_img_size': 640,
        'cls_model_path': r"D:\DATA\20250519RENBAO\trainV8Pose_closePeople\models\yolov8sCls_320_closePeople_20250605_165609_cls_format_data__frontCamImg0_100_cls4_augMean3_padSquare\weights\best.pt",
        'cls_model_type': cls_model_type,
        # v3
        'save_dir': rf"{root_path}/temp/1749118362.3255725_frontCamImg0_100_cls4_augMean3_padSquare",
        'img_glob': r'D:\DATA\20250519RENBAO\trainV8Pose_closePeople\caitu\64\frontCam_testdata200_t23721-2025-06-04_09-57-50.mp4_frames\*.jpg',
        'action_celue': 'v3',
        'is_generate_labelmepose': True,
        'is_save_img': True,

        # 'is_compare_annotation': True,
        # 'true_annotation_dir': r'D:\DATA\20250519RENBAO\caitu\pose_test_dataset250-labelme',
        # 'pre_annotation_dir': r'D:\DATA\20250519RENBAO\temp\annotated',

    }
    compare_annotation(test_config)


# save_dir D:\WeChat\xwechat_files\wxid_ig5xjsfld7cp22_08e6\msg\file\2025-06\JPEGImages2\JPEGImages2_result
# cls_model_type img100_200_cls7_augmean3_pad
# box_error_count 21
# baggage_place acc: 475 / 631 0.75
# none_action acc: 330 / 569 0.58
# acc: 805 / 1200 0.67